#!/usr/bin/env python3
"""
数据格式转换智能体 V2
使用LangGraph的create_react_agent重构
"""

import json
import re
from typing import Dict, List, Any, Optional
from typing_extensions import TypedDict
from dataclasses import dataclass
from langchain.schema import BaseMessage, HumanMessage, SystemMessage
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import InMemorySaver
from langchain.tools import tool
from config import DS_CONFIG, QWEN_CONFIG

# 导入LLM配置
try:
    from langchain_deepseek import ChatDeepSeek
    from langchain_openai import ChatOpenAI
except ImportError:
    print("⚠️ 缺少langchain依赖，请安装: pip install langchain-deepseek langchain-openai")

# 状态定义
class FormatterState(TypedDict):
    """格式转换智能体的状态"""
    user_request: str
    query_results: List[Dict[str, Any]]
    summary: str
    chart_type: str
    echarts_data: Optional[Dict[str, Any]]
    final_output: Optional[Dict[str, Any]]
    messages: List[BaseMessage]
    error: Optional[str]

@dataclass
class FormatterInput:
    """格式转换智能体的输入数据"""
    user_request: str
    query_results: List[Dict[str, Any]]
    summary: str

class DataFormatterAgentV2:
    """数据格式转换智能体 V2 - 使用LangGraph"""
    
    def __init__(self, config=DS_CONFIG):
        """初始化格式转换智能体"""
        self.config = config
        self.llm = self._init_llm()
        self.graph = self._create_graph()
        
    def _init_llm(self):
        """初始化LLM模型"""
        try:
            if self.config.MODEL == 'deepseek-v3':
                return ChatDeepSeek(
                    model="deepseek-v3",
                    base_url=self.config.BASE_URL,
                    api_key=self.config.ALIYUN_API_KEY,
                    api_base=self.config.BASE_URL,
                    temperature=0.3,
                    streaming=False
                )
            elif self.config.MODEL == 'qwen-max':
                return ChatOpenAI(
                    model="qwen-max",
                    base_url=self.config.BASE_URL,
                    api_key=self.config.ALIYUN_API_KEY,
                    temperature=0.3,
                    streaming=False
                )
            else:
                print(f"⚠️ 不支持的模型: {self.config.MODEL}")
                return None
        except Exception as e:
            print(f"❌ LLM初始化失败: {e}")
            return None
    
    def _create_graph(self):
        """创建LangGraph工作流"""
        if not self.llm:
            print("❌ LLM未初始化，无法创建图")
            return None
        
        # 创建图
        workflow = StateGraph(FormatterState)
        
        # 添加节点
        workflow.add_node("analyze", self._analyze_node)
        workflow.add_node("select", self._select_node)
        workflow.add_node("generate", self._generate_node)
        workflow.add_node("validate", self._validate_node)
        
        # 设置入口点
        workflow.set_entry_point("analyze")
        
        # 添加边
        workflow.add_edge("analyze", "select")
        workflow.add_edge("select", "generate")
        workflow.add_edge("generate", "validate")
        workflow.add_edge("validate", END)
        
        # 添加检查点
        checkpointer = InMemorySaver()
        
        return workflow.compile(checkpointer=checkpointer)
    
    def _analyze_node(self, state: FormatterState) -> FormatterState:
        """分析数据结构的节点"""
        try:
            result = self._analyze_data_structure(state)
            print(f"🔍 数据结构分析: {result}")
        except Exception as e:
            print(f"❌ 数据结构分析失败: {e}")
        return state
    
    def _select_node(self, state: FormatterState) -> FormatterState:
        """选择图表类型的节点"""
        try:
            result = self._select_chart_type(state)
            print(f"🎯 图表类型选择: {result}")
        except Exception as e:
            print(f"❌ 图表类型选择失败: {e}")
        return state
    
    def _generate_node(self, state: FormatterState) -> FormatterState:
        """生成ECharts数据的节点"""
        try:
            result = self._generate_echarts_data(state)
            print(f"📊 ECharts数据生成: {result}")
        except Exception as e:
            print(f"❌ ECharts数据生成失败: {e}")
        return state
    
    def _validate_node(self, state: FormatterState) -> FormatterState:
        """验证输出格式的节点"""
        try:
            result = self._validate_output(state)
            print(f"✅ 输出格式验证: {result}")
        except Exception as e:
            print(f"❌ 输出格式验证失败: {e}")
        return state
    
    def _analyze_data_structure(self, state: FormatterState) -> str:
        """分析数据结构"""
        try:
            query_results = state["query_results"]
            if not query_results:
                return "数据结构分析：查询结果为空"
            
            first_row = query_results[0]
            fields = list(first_row.keys())
            
            # 分析字段类型
            numeric_fields = []
            text_fields = []
            for field, value in first_row.items():
                if isinstance(value, (int, float)):
                    numeric_fields.append(field)
                else:
                    text_fields.append(field)
            
            analysis = {
                "total_records": len(query_results),
                "fields": fields,
                "numeric_fields": numeric_fields,
                "text_fields": text_fields,
                "sample_data": query_results[:3]
            }
            
            return f"数据结构分析完成：{json.dumps(analysis, ensure_ascii=False)}"
            
        except Exception as e:
            return f"数据结构分析失败：{str(e)}"
    
    def _select_chart_type(self, state: FormatterState) -> str:
        """选择图表类型"""
        try:
            user_request = state["user_request"].lower()
            query_results = state["query_results"]
            
            # 用户明确要求
            if "饼图" in user_request or "占比" in user_request:
                chart_type = "pie"
            elif "折线图" in user_request or "趋势" in user_request:
                chart_type = "line"
            elif "散点图" in user_request or "关系" in user_request:
                chart_type = "scatter"
            else:
                # 数据特征分析
                if not query_results:
                    chart_type = "bar"
                else:
                    first_row = query_results[0]
                    numeric_fields = [k for k, v in first_row.items() if isinstance(v, (int, float))]
                    
                    if len(numeric_fields) >= 2:
                        chart_type = "scatter"
                    elif len(query_results) <= 8 and any(keyword in user_request for keyword in ["占比", "分布", "比例"]):
                        chart_type = "pie"
                    elif any(keyword in user_request for keyword in ["趋势", "变化", "时间", "月份"]):
                        chart_type = "line"
                    else:
                        chart_type = "bar"
            
            # 更新状态
            state["chart_type"] = chart_type
            return f"图表类型选择：{chart_type}"
            
        except Exception as e:
            return f"图表类型选择失败：{str(e)}"
    
    def _generate_echarts_data(self, state: FormatterState) -> str:
        """生成ECharts数据"""
        try:
            query_results = state["query_results"]
            chart_type = state["chart_type"]
            
            if not query_results:
                state["echarts_data"] = None
                return "ECharts数据生成：无数据"
            
            first_row = query_results[0]
            fields = list(first_row.keys())
            
            if chart_type == "bar":
                echarts_data = self._generate_bar_data(query_results, fields)
            elif chart_type == "pie":
                echarts_data = self._generate_pie_data(query_results, fields)
            elif chart_type == "line":
                echarts_data = self._generate_line_data(query_results, fields)
            elif chart_type == "scatter":
                echarts_data = self._generate_scatter_data(query_results, fields)
            else:
                echarts_data = self._generate_bar_data(query_results, fields)
            
            state["echarts_data"] = echarts_data
            return f"ECharts数据生成完成：{chart_type}格式"
            
        except Exception as e:
            state["echarts_data"] = None
            return f"ECharts数据生成失败：{str(e)}"
    
    def _validate_output(self, state: FormatterState) -> str:
        """验证输出格式"""
        try:
            # 构建最终输出
            final_output = {
                "summary": state["summary"],
                "chart_type": state["chart_type"],
                "data": state["echarts_data"]
            }
            
            # 验证必需字段
            if not final_output["summary"]:
                final_output["summary"] = "数据格式转换完成"
            
            if not final_output["chart_type"]:
                final_output["chart_type"] = "bar"
            
            state["final_output"] = final_output
            return "输出格式验证完成"
            
        except Exception as e:
            state["error"] = str(e)
            return f"输出格式验证失败：{str(e)}"
    
    def _generate_bar_data(self, query_results: List[Dict[str, Any]], fields: List[str]) -> Dict[str, Any]:
        """生成柱状图数据"""
        categories = [str(row.get(fields[0], "")) for row in query_results]
        values = [row.get(fields[1], 0) if len(fields) > 1 else 0 for row in query_results]
        
        return {
            "xAxis": {
                "type": "category",
                "data": categories
            },
            "yAxis": {
                "type": "value",
                "name": fields[1] if len(fields) > 1 else "数值"
            },
            "series": [{
                "name": fields[1] if len(fields) > 1 else "数值",
                "type": "bar",
                "data": values,
                "itemStyle": {"color": "#5470c6"}
            }]
        }
    
    def _generate_pie_data(self, query_results: List[Dict[str, Any]], fields: List[str]) -> Dict[str, Any]:
        """生成饼图数据"""
        data = []
        for row in query_results:
            name = str(row.get(fields[0], ""))
            value = row.get(fields[1], 0) if len(fields) > 1 else 0
            data.append({"name": name, "value": value})
        
        return {
            "series": [{
                "name": fields[1] if len(fields) > 1 else "数值",
                "type": "pie",
                "radius": "50%",
                "data": data
            }]
        }
    
    def _generate_line_data(self, query_results: List[Dict[str, Any]], fields: List[str]) -> Dict[str, Any]:
        """生成折线图数据"""
        categories = [str(row.get(fields[0], "")) for row in query_results]
        values = [row.get(fields[1], 0) if len(fields) > 1 else 0 for row in query_results]
        
        return {
            "xAxis": {
                "type": "category",
                "data": categories
            },
            "yAxis": {
                "type": "value",
                "name": fields[1] if len(fields) > 1 else "数值"
            },
            "series": [{
                "name": fields[1] if len(fields) > 1 else "数值",
                "type": "line",
                "data": values,
                "smooth": True
            }]
        }
    
    def _generate_scatter_data(self, query_results: List[Dict[str, Any]], fields: List[str]) -> Dict[str, Any]:
        """生成散点图数据"""
        data = []
        for row in query_results:
            x = row.get(fields[0], 0)
            y = row.get(fields[1], 0) if len(fields) > 1 else 0
            data.append([x, y])
        
        return {
            "xAxis": {
                "type": "value",
                "name": fields[0]
            },
            "yAxis": {
                "type": "value",
                "name": fields[1] if len(fields) > 1 else "数值"
            },
            "series": [{
                "name": "数据点",
                "type": "scatter",
                "data": data,
                "symbolSize": 10
            }]
        }
    
    def format_data(self, formatter_input: FormatterInput) -> Dict[str, Any]:
        """
        使用LangGraph进行数据格式转换
        
        Args:
            formatter_input: 输入数据
            
        Returns:
            Dict: 标准格式输出
        """
        try:
            print("🔄 开始LangGraph数据格式转换...")
            
            if not self.graph:
                print("❌ 图未初始化，使用备用逻辑")
                return self._fallback_format(formatter_input)
            
            # 构建初始状态
            initial_state = FormatterState(
                user_request=formatter_input.user_request,
                query_results=formatter_input.query_results,
                summary=formatter_input.summary,
                chart_type="",
                echarts_data=None,
                final_output=None,
                messages=[
                    SystemMessage(content=self._get_system_prompt()),
                    HumanMessage(content=self._build_user_prompt(formatter_input))
                ],
                error=None
            )
            
            # 执行图
            result = self.graph.invoke(initial_state, {"configurable": {"thread_id": "formatter_thread"}})
            
            # 获取最终输出
            if result.get("final_output"):
                print("✅ LangGraph格式转换成功")
                return result["final_output"]
            else:
                print("⚠️ LangGraph转换失败，使用备用逻辑")
                return self._fallback_format(formatter_input)
                
        except Exception as e:
            print(f"❌ LangGraph执行失败: {e}")
            return self._fallback_format(formatter_input)
    
    def _get_system_prompt(self) -> str:
        """获取系统prompt"""
        return """你是一个专业的数据格式转换智能体，负责将数据分析结果转换为标准的前端可视化格式。

## 你的职责
1. 分析用户请求和数据特征
2. 选择最合适的图表类型
3. 生成符合ECharts标准的配置数据
4. 输出标准格式的结果

## 输出格式要求
你必须输出包含以下三个字段的JSON格式：
1. "summary": 用户友好的数据总结
2. "chart_type": 选择的图表类型（bar/pie/line/scatter）
3. "data": ECharts配置数据

## 图表类型选择规则
- bar: 柱状图 - 适用于分类数据对比
- pie: 饼图 - 适用于占比分析（≤8个类别）
- line: 折线图 - 适用于趋势分析
- scatter: 散点图 - 适用于相关性分析

## 处理步骤
1. 使用_analyze_data_structure工具分析数据结构
2. 使用_select_chart_type工具选择图表类型
3. 使用_generate_echarts_data工具生成ECharts数据
4. 使用_validate_output工具验证输出格式

请严格按照上述步骤处理，确保输出格式正确。"""
    
    def _build_user_prompt(self, formatter_input: FormatterInput) -> str:
        """构建用户prompt"""
        return f"""请将以下数据分析结果转换为标准格式：

用户请求: {formatter_input.user_request}
数据行数: {len(formatter_input.query_results)}
现有总结: {formatter_input.summary}

请按照系统提示的步骤进行处理，最终输出标准格式的JSON结果。"""
    
    def _fallback_format(self, formatter_input: FormatterInput) -> Dict[str, Any]:
        """备用格式转换逻辑"""
        try:
            print("🔄 使用备用格式转换逻辑...")
            
            # 简单的图表类型选择
            user_request_lower = formatter_input.user_request.lower()
            if "饼图" in user_request_lower or "占比" in user_request_lower:
                chart_type = "pie"
            elif "折线图" in user_request_lower or "趋势" in user_request_lower:
                chart_type = "line"
            elif "散点图" in user_request_lower or "关系" in user_request_lower:
                chart_type = "scatter"
            else:
                chart_type = "bar"
            
            # 简单的ECharts数据生成
            if formatter_input.query_results:
                first_row = formatter_input.query_results[0]
                fields = list(first_row.keys())
                
                if chart_type == "bar":
                    echarts_data = self._generate_bar_data(formatter_input.query_results, fields)
                elif chart_type == "pie":
                    echarts_data = self._generate_pie_data(formatter_input.query_results, fields)
                elif chart_type == "line":
                    echarts_data = self._generate_line_data(formatter_input.query_results, fields)
                elif chart_type == "scatter":
                    echarts_data = self._generate_scatter_data(formatter_input.query_results, fields)
                else:
                    echarts_data = self._generate_bar_data(formatter_input.query_results, fields)
            else:
                echarts_data = None
            
            return {
                "summary": formatter_input.summary,
                "chart_type": chart_type,
                "data": echarts_data
            }
            
        except Exception as e:
            print(f"❌ 备用格式转换失败: {e}")
            return {
                "summary": formatter_input.summary,
                "chart_type": "bar",
                "data": None
            }

# 全局实例
_formatter_agent_v2 = None

def get_formatter_agent_v2(config=DS_CONFIG) -> DataFormatterAgentV2:
    """获取格式转换智能体V2实例"""
    global _formatter_agent_v2
    if _formatter_agent_v2 is None:
        _formatter_agent_v2 = DataFormatterAgentV2(config)
    return _formatter_agent_v2

def format_data_to_standard_v2(user_request: str, query_results: List[Dict[str, Any]], summary: str, config=DS_CONFIG) -> Dict[str, Any]:
    """
    使用LangGraph将数据分析结果转换为标准格式
    
    Args:
        user_request: 用户原始请求
        query_results: 数据库查询结果
        summary: 现有智能体的分析总结
        config: LLM配置
        
    Returns:
        Dict: 包含summary、chart_type和data的标准格式
    """
    formatter = get_formatter_agent_v2(config)
    
    formatter_input = FormatterInput(
        user_request=user_request,
        query_results=query_results,
        summary=summary
    )
    
    return formatter.format_data(formatter_input)

if __name__ == "__main__":
    # 测试代码
    test_data = [
        {"ATTR_CITY": "朝阳区", "count": 185},
        {"ATTR_CITY": "丰台区", "count": 103},
        {"ATTR_CITY": "东城区", "count": 97}
    ]
    
    result = format_data_to_standard_v2(
        user_request="请统计北京不同区中高端用户数",
        query_results=test_data,
        summary="北京市各区中高端用户数量统计完成"
    )
    
    print("测试结果:")
    print(json.dumps(result, ensure_ascii=False, indent=2))
